from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import cv2
import os
import uuid
import shutil
import base64
from zhongfeng import lipClassification

app = FastAPI()

# 目录路径
UPLOAD_DIR = "/home/mengxiaowang/fuxian/lzy/file4HW/image_to_test"
OUTPUT_DIR = "/home/mengxiaowang/fuxian/lzy/file4HW/image_to_getback"

# 创建目录（如果不存在）
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)


@app.post("/process_image/")
async def process_image(file: UploadFile = File(...)):
    # 生成唯一文件名
    file_extension = file.filename.split(".")[-1]
    file_name = f"{uuid.uuid4()}.{file_extension}"
    file_path = os.path.join(UPLOAD_DIR, file_name)

    # 保存上传的文件
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # 调用图像处理函数获取结果
    result, processed_image = lipClassification(file_path)

    # 将处理后的图像保存到输出目录
    output_image_path = os.path.join(OUTPUT_DIR, file_name)
    cv2.imwrite(output_image_path, processed_image)

    # 将图像转换为base64格式以便在响应中发送
    _, buffer = cv2.imencode('.jpg', processed_image)
    img_str = base64.b64encode(buffer).decode('utf-8')

    return JSONResponse(content={
        "result": result,
        "image": img_str
    })


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
